import os
import random
import librosa
import soundfile as sf
import numpy as np
from pathlib import Path
from typing import List, Optional, Dict


class SampleGenerator:
    """负责从源音乐库中批量生成测试样本，并可选择性添加噪声。"""
    def __init__(self, input_dir: Path, output_dir: Path, config: Dict):
        self.input_dir = Path(input_dir)
        self.output_dir = Path(output_dir)
        self.min_length = 8
        self.max_length = 20
        self.noise_config = config.get('batch_test', {}).get('noise_generation', {})
        self.default_snr_db = self.noise_config.get('default_snr_db', 20)
        self.output_dir.mkdir(exist_ok=True, parents=True)

    def _add_noise(self, y: np.ndarray, snr_db: int) -> np.ndarray:
        """根据信噪比(SNR)向音频信号添加高斯白噪声。"""
        signal_power = np.mean(y**2)
        if signal_power == 0: return y # 避免除以零
        noise_power = signal_power / (10**(snr_db / 10))
        noise = np.random.normal(0, np.sqrt(noise_power), len(y))
        return y + noise

    def _cut_single_file(self, input_path: Path, add_noise: bool, snr_db: int) -> Optional[str]:
        """对单个文件进行裁剪，并根据参数添加噪声。"""
        try:
            y, sr = librosa.load(input_path, sr=None, mono=True)
            duration = librosa.get_duration(y=y, sr=sr)
            if duration < self.min_length: return None

            max_possible_length = min(self.max_length, int(duration))
            if self.min_length > max_possible_length: return None

            clip_length = random.randint(self.min_length, max_possible_length)
            max_start = int(duration) - clip_length
            if max_start < 0: return None
            start_seconds = random.randint(0, max_start)

            start_idx = int(start_seconds * sr)
            end_idx = int((start_seconds + clip_length) * sr)
            clipped_y = y[start_idx:end_idx]

            # 使用规范化的命名格式
            filename_base = f"{input_path.stem}---{start_seconds}s---{clip_length}s"
            if add_noise:
                clipped_y = self._add_noise(clipped_y, snr_db)
                output_file = self.output_dir / f"{filename_base}---noise{snr_db}db.wav"
            else:
                output_file = self.output_dir / f"{filename_base}.wav"

            sf.write(file=output_file, data=clipped_y, samplerate=sr, subtype='PCM_16')
            return str(output_file)
        except Exception as e:
            print(f"Error cutting {input_path.name}: {e}")
            return None

    def generate_for_all(self, add_noise: bool = False, snr_db: Optional[int] = None) -> List[str]:
        """为输入目录中的所有音频生成一个测试样本。"""
        snr = snr_db if snr_db is not None else self.default_snr_db
        generated_files = []
        audio_files = list(self.input_dir.glob("*.wav")) + list(self.input_dir.glob("*.mp3"))
        for audio_file in audio_files:
            if audio_file.is_file():
                result = self._cut_single_file(audio_file, add_noise, snr)
                if result:
                    generated_files.append(Path(result).name)
        return generated_files
